# 导入必要的库
from sklearn.datasets import load_digits  # 用于加载手写数字数据集
from sklearn.neighbors import KNeighborsClassifier  # 用于创建 KNN 模型
from sklearn.model_selection import train_test_split  # 用于分割数据集
from tqdm import tqdm
import pickle
import matplotlib.pyplot as plt

train_dict = {}

knn_list = []

def return_model(num):
    digits = load_digits()  # 加载手写数字数据集
    X = digits.data  # 获取样本数据
    y = digits.target  # 获取样本标签

    # 分割数据集为训练集和测试集
    # test_size=0.2 表示测试集占总数据集的 20%
    # random_state=42 表示随机种子，保证每次运行代码时分割的结果都相同，这主要是为了方便大家调试代码
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

    # 创建 KNN 模型并进行训练
    knn = KNeighborsClassifier(n_neighbors=num)  # 创建 KNN 模型
    knn.fit(X_train, y_train)  # 使用训练集进行训练

    return knn


def train(num):
    digits = load_digits()  # 加载手写数字数据集
    X = digits.data  # 获取样本数据
    y = digits.target  # 获取样本标签

    # 分割数据集为训练集和测试集
    # test_size=0.2 表示测试集占总数据集的 20%
    # random_state=42 表示随机种子，保证每次运行代码时分割的结果都相同，这主要是为了方便大家调试代码
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

    # 创建 KNN 模型并进行训练
    knn = KNeighborsClassifier(n_neighbors=num)  # 创建 KNN 模型
    knn.fit(X_train, y_train)  # 使用训练集进行训练
    knn_list.append(knn)
    # 直接计算模型的预测准确率，利用测试集进行测试
    acc = knn.score(X_test, y_test)  # 计算预测准确率

    # print(f'预测准确率为：{acc * 100}%')  # 打印预测准确率，查看预测准确率
    return acc * 100


def return_key(val):
    for ky, vue in train_dict.items():
        if vue == val:
            return ky


for i in tqdm(range(1, 41), unit="单位"):
    acc = train(i)
    train_dict[i] = acc

print(f'当k值为{return_key(max(train_dict.values()))}时,对应的准确率最高为：{max(train_dict.values())}%')
key_list = []
value_list = []
for key, value in train_dict.items():
    key_list.append(key)
    value_list.append(value)
plt.plot(key_list, value_list, linestyle='-', color='blue')
plt.title('Accuracy of different k values')
plt.xlabel('k values')
plt.ylabel('Accuracy')
plt.text(return_key(max(train_dict.values())), max(train_dict.values()),
         f'k={return_key(max(train_dict.values()))},Accuracy={max(train_dict.values())}', fontsize=12, color='red')
plt.axvline(return_key(max(train_dict.values())), color='r')
plt.show()
# plt.savefig("accuracy_plot.pdf")
pkl = return_model(return_key(max(train_dict.values())))
with open("best_knn_model.pkl", 'wb') as f:
    pickle.dump(pkl, f)
